import sqlite3
import pandas as pd
from contextlib import contextmanager

class SQLiteMemory:
    def __init__(self, dbpath):  # dbpath数据库文件路径
        self.dbpath = dbpath

    @contextmanager
    def connect(self):
        conn = sqlite3.connect(self.dbpath)
        try:
            yield conn
        finally:
            conn.close()

    # 初始化内存数据库，把DataFrame插入到内存数据库中
    def initialize_memory_database(self, df, table_name, if_exists='append'):
        with self.connect() as conn:
            try:
                # 将新数据插入到内存数据库的表中
                df.to_sql(table_name, con=conn, if_exists=if_exists, index=False)
                # 提交内存数据库的更改
                conn.commit()
                print("ok")
            except Exception as e:
                print(f"An error occurred: {e}")

    # 使用事务批量删除
    def delete_from_sqlite_with_transaction(self, df, table_name, column):
        """
        使用事务批量删除SQLite数据库中的记录。

        参数：
        df: DataFrame对象，包含要删除的记录的标识。
        table_name: 数据库中的表名。
        column: 要匹配的列名。
        """
        with self.connect() as conn:
            cursor = conn.cursor()
            # 开始一个新的事务
            cursor.execute("BEGIN TRANSACTION;")
            try:
                # 遍历DataFrame对象，删除每一行对应的记录
                for index, row in df.iterrows():
                    # 构造SQL语句
                    sql_delete = f"DELETE FROM {table_name} WHERE {column} = ?"
                    # 执行SQL语句
                    cursor.execute(sql_delete, (row[column],))
                # 提交事务
                conn.commit()
            except Exception as e:
                print(f"An error occurred: {e}")
                # 回滚事务
                conn.rollback()

    # 查询内存数据库中的数据，如果查询成功，返回DataFrame，否则返回None
    def query(self, sql):
        with self.connect() as conn:
            return pd.read_sql_query(sql, conn)
